import os
import numpy as np
from math import dist
import time
import logging
import imageio
import matplotlib.pyplot as plt
import io
import networkx as nx

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class GraphManager:
    def __init__(self):
        self.node_coordinates = []  # 节点坐标列表
        self.adj_matrix = []       # 邻接矩阵
        self.dist_3d = []          # 三维距离矩阵
        self.f_matrix = []         # Floyd-Warshall 最短路径矩阵
        self.node_num = 0
        # 加载图数据
        self.load_graph_data()
    
    def load_graph_data(self):
        """加载图数据（邻接矩阵和预计算的距离矩阵）"""
        root = '../datas/'
        try:
            # 加载邻接矩阵
            adj_matrix_path = os.path.join(root, 'adj_matrix.txt')
            if not os.path.exists(adj_matrix_path):
                raise FileNotFoundError(f"邻接矩阵文件不存在: {adj_matrix_path}")
            with open(adj_matrix_path, 'r') as f:
                self.node_num = int(f.readline().strip())
                self.adj_matrix = np.zeros((self.node_num, self.node_num), dtype=int)
                for i in range(self.node_num):
                    row = list(map(int, f.readline().split()))
                    self.adj_matrix[i] = row
            # 加载节点坐标
            points_path = os.path.join(root, 'points.txt')
            self.load_node_coordinates(points_path)
            # 加载预计算的距离矩阵 - 仅使用NPY格式
            dist_npy_path = os.path.join(root, 'D_bin.npy')
            if not os.path.exists(dist_npy_path):
                raise FileNotFoundError(f"距离矩阵文件不存在: {dist_npy_path}")
            # 加载NPY格式
            start_time = time.time()
            self.dist_3d = np.load(dist_npy_path)
            load_time = time.time() - start_time
            # 验证矩阵形状
            if self.dist_3d.shape != (self.node_num+1, self.node_num+1, self.node_num+1):
                raise ValueError(f"距离矩阵形状不匹配: 预期({self.node_num},{self.node_num},{self.node_num}), "
                                 f"实际{self.dist_3d.shape}")
            logger.info(f"从NPY文件加载距离矩阵成功! 节点数: {self.node_num}, 耗时: {load_time:.4f}秒")
            # 计算最短路径矩阵 (Floyd-Warshall)
            logger.info("开始计算最短路径矩阵...")
            self.compute_shortest_paths()
        except Exception as e:
            logger.error(f"加载图数据错误: {e}")
            raise  # 重新抛出异常以便上层处理
    
    def load_node_coordinates(self, filename):
        """从文件加载节点坐标，文件不存在时抛出异常"""
        if not os.path.exists(filename):
            raise FileNotFoundError(f"节点坐标文件不存在: {filename}")
        try:
            self.node_coordinates = []
            with open(filename, 'r') as f:
                for line in f:
                    # 尝试两种格式：逗号分隔或空格分隔
                    if ',' in line:
                        parts = line.strip().split(',')
                    else:
                        parts = line.strip().split()
                    if len(parts) >= 2:
                        x = float(parts[0])
                        y = float(parts[1])
                        self.node_coordinates.append((x, y))
            # 验证节点数量是否匹配
            if len(self.node_coordinates) != self.node_num:
                raise ValueError(f"坐标文件节点数不匹配: 坐标文件 {len(self.node_coordinates)} 个节点, "
                                 f"邻接矩阵 {self.node_num} 个节点")
            logger.info(f"成功加载 {len(self.node_coordinates)} 个节点的坐标")
            return True
        except Exception as e:
            logger.error(f"加载节点坐标错误: {e}")
            raise  # 重新抛出异常
    
    def compute_shortest_paths(self):
        """
        使用 NetworkX 的 floyd_warshall_numpy 加速最短路径矩阵计算
        """
        # 1. 从邻接矩阵构建有向图
        G = nx.from_numpy_array(self.adj_matrix, create_using=nx.DiGraph())

        # 2. 使用 NetworkX 的 floyd_warshall_numpy 计算所有节点对的最短路径
        dist_matrix = nx.floyd_warshall_numpy(G)

        # 3. 定义不可达距离（原代码中为 255）
        INF = 255

        # 4. 替换不可达路径为 INF，并四舍五入为整数
        dist_matrix[np.isinf(dist_matrix)] = INF
        self.f_matrix = np.rint(dist_matrix).astype(int)

        # 5. 确保对角线为 0（自己到自己的距离）
        np.fill_diagonal(self.f_matrix, 0)

        logger.info("使用 NetworkX 成功计算最短路径矩阵")
    def position_to_node(self, x, y):
        """将实际坐标映射到最近的节点"""
        min_dist = float('inf')
        closest_node = 0
        for idx, (nx, ny) in enumerate(self.node_coordinates):
            d = dist((x, y), (nx, ny))
            if d < min_dist:
                min_dist = d
                closest_node = idx
        return closest_node
    
    def node_to_position(self, node_id):
        """将节点ID转换为实际坐标"""
        if 0 <= node_id < len(self.node_coordinates):
            return self.node_coordinates[node_id]
        return (0.0, 0.0)  # 默认位置
    
    def get_neighbors(self, node_id):
        """获取节点的邻居（不包括自身）"""
        neighbors = []
        for j in range(self.node_num):
            if self.adj_matrix[node_id][j] == 1:
                neighbors.append(j)
        return neighbors

class GameStrategy:
    def __init__(self, graph):
        self.graph = graph
        self.capture_count = 0

    def calculate_next_move(self, evader_node, pursuer1_node, pursuer2_node):
        if self.is_captured(evader_node, pursuer1_node, pursuer2_node):
            return evader_node, pursuer1_node, pursuer2_node

        evader_neighbors = self.graph.get_neighbors(evader_node)
        p1_neighbors = self.graph.get_neighbors(pursuer1_node)
        p2_neighbors = self.graph.get_neighbors(pursuer2_node)
        
        # 逃脱者决策：最大化最小距离
        best_evader = evader_node
        max_min_value = -1
        best_evader_dist_sum = float('-inf')  # 距离和初始值
        
        for e_next in evader_neighbors:
            min_value = float('inf')
            for p1_next in p1_neighbors:
                for p2_next in p2_neighbors:
                    value = self.graph.dist_3d[p1_next+1][p2_next+1][e_next+1]
                    if value < min_value:
                        min_value = value
            
            # 平局处理：计算路径距离和
            dist_sum = (self.graph.f_matrix[pursuer1_node][e_next] + 
                        self.graph.f_matrix[pursuer2_node][e_next])
            
            if min_value > max_min_value:
                max_min_value = min_value
                best_evader = e_next
                best_evader_dist_sum = dist_sum
            elif min_value == max_min_value and dist_sum > best_evader_dist_sum:
                best_evader = e_next
                best_evader_dist_sum = dist_sum

        # 追捕者决策：最小化最大距离
        best_p1 = pursuer1_node
        best_p2 = pursuer2_node
        min_max_value = float('inf')
        best_pursuer_dist_sum = float('inf')  # 距离和初始值
        
        for p1_next in p1_neighbors:
            for p2_next in p2_neighbors:
                max_value = -1
                for e_next in evader_neighbors:
                    value = self.graph.dist_3d[p1_next+1][p2_next+1][e_next+1]
                    if value > max_value:
                        max_value = value
                
                # 平局处理：计算路径距离和
                dist_sum = (self.graph.f_matrix[p1_next][evader_node] + 
                            self.graph.f_matrix[p2_next][evader_node])
                
                if max_value < min_max_value:
                    min_max_value = max_value
                    best_p1 = p1_next
                    best_p2 = p2_next
                    best_pursuer_dist_sum = dist_sum
                elif (max_value == min_max_value and 
                      dist_sum < best_pursuer_dist_sum):
                    best_p1 = p1_next
                    best_p2 = p2_next
                    best_pursuer_dist_sum = dist_sum

        return best_evader, best_p1, best_p2
    def is_captured(self, evader_node, pursuer1_node, pursuer2_node):
        """检查逃避者是否被捕获"""
        # 条件1: 逃避者与任一追捕者在同一节点
        if evader_node == pursuer1_node or evader_node == pursuer2_node:
            return True
        # 条件2: 逃避者与任一追捕者直接相连
        if (self.graph.adj_matrix[evader_node][pursuer1_node] == 1 or 
            self.graph.adj_matrix[evader_node][pursuer2_node] == 1):
            return True
        return False


def create_frame(graph, evader_node, pursuer1_node, pursuer2_node):
    """创建单帧图像，可视化当前状态"""
    fig, ax = plt.subplots(figsize=(8, 8))
    
    # 绘制边（仅绘制上三角矩阵避免重复）
    for i in range(graph.node_num):
        for j in range(i+1, graph.node_num):
            if graph.adj_matrix[i][j] == 1:
                ax.plot(
                    [graph.node_coordinates[i][0], graph.node_coordinates[j][0]],
                    [graph.node_coordinates[i][1], graph.node_coordinates[j][1]],
                    'k-', alpha=0.3
                )
    
    # 绘制所有节点
    all_x = [coord[0] for coord in graph.node_coordinates]
    all_y = [coord[1] for coord in graph.node_coordinates]
    ax.scatter(all_x, all_y, c='black', s=20)
    
    # 绘制角色位置
    ax.scatter(*graph.node_coordinates[evader_node], c='red', s=100, marker='o', label='Evader')
    ax.scatter(*graph.node_coordinates[pursuer1_node], c='blue', s=100, marker='^', label='Pursuer 1')
    ax.scatter(*graph.node_coordinates[pursuer2_node], c='green', s=100, marker='s', label='Pursuer 2')
    
    ax.legend(loc='upper right')
    ax.set_title('Evader vs Pursuers')
    
    # 自动计算绘图范围并添加 10% 边距
    if graph.node_coordinates:
        x_coords = [coord[0] for coord in graph.node_coordinates]
        y_coords = [coord[1] for coord in graph.node_coordinates]
        
        x_min, x_max = min(x_coords), max(x_coords)
        y_min, y_max = min(y_coords), max(y_coords)
        
        x_range = x_max - x_min
        y_range = y_max - y_min
        
        # 处理零范围情况（例如所有点在同一坐标）
        x_margin = x_range * 0.1 if x_range > 0 else 0.1
        y_margin = y_range * 0.1 if y_range > 0 else 0.1
        
        ax.set_xlim(x_min - x_margin, x_max + x_margin)
        ax.set_ylim(y_min - y_margin, y_max + y_margin)
    
    ax.set_aspect('equal')
    ax.axis('off')
    
    # 将图像保存到内存中的字节流
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight')
    plt.close(fig)
    buf.seek(0)
    image = imageio.imread(buf)
    buf.close()
    return image



if __name__ == "__main__":
    outputpath = "../datas/simulate_output.gif"
    graph = GraphManager()
    strategy = GameStrategy(graph)
    
    # 初始化随机位置
    evader_node = np.random.randint(0, graph.node_num)
    pursuer1_node = np.random.randint(0, graph.node_num)
    pursuer2_node = np.random.randint(0, graph.node_num)
    
    frames = []
    max_steps = 100  # 最大步数限制
    while True:
        # 检查是否结束
        if strategy.is_captured(evader_node, pursuer1_node, pursuer2_node):
            logger.info("逃避者已被捕获!")
            break
        
        # 生成当前帧并添加到frames
        frame = create_frame(graph, evader_node, pursuer1_node, pursuer2_node)
        frames.append(frame)
        
        # 检查是否被捕获
        if strategy.is_captured(evader_node, pursuer1_node, pursuer2_node):
            break
        
        # 计算下一步
        next_evader, next_p1, next_p2 = strategy.calculate_next_move(
            evader_node, pursuer1_node, pursuer2_node
        )
        
        if next_evader == evader_node and next_p1 == pursuer1_node and next_p2 == pursuer2_node:
            logger.info("逃避者和追捕者没有移动！")
            break
        
        evader_node, pursuer1_node, pursuer2_node = next_evader, next_p1, next_p2
        
        
        logger.info(f"剩余步骤：{max_steps}，当前状态: 逃避者 {evader_node}, 追捕者1 {pursuer1_node}, 追捕者2 {pursuer2_node}")
        
        
        max_steps -= 1
        if max_steps <= 0:
            break
    
    # 保存为 GIF
    imageio.mimsave(outputpath, frames, duration=0.5)
    print(f"GIF 已保存到 {outputpath}")